#include <stdio.h>
#include <mutex>
#include "rknn_api.h"



#include "opencv2/core/core.hpp"
#include "opencv2/highgui/highgui.hpp"
#include "opencv2/imgproc/imgproc.hpp"

#include "coreNum.hpp"
#include "rkYolov8_seg.hpp"
#include <QtDebug>

unsigned char class_colors[][3] = {
    {255, 56, 56},   // 'FF3838'
    {255, 157, 151}, // 'FF9D97'
    {255, 112, 31},  // 'FF701F'
    {255, 178, 29},  // 'FFB21D'
    {207, 210, 49},  // 'CFD231'
    {72, 249, 10},   // '48F90A'
    {146, 204, 23},  // '92CC17'
    {61, 219, 134},  // '3DDB86'
    {26, 147, 52},   // '1A9334'
    {0, 212, 187},   // '00D4BB'
    {44, 153, 168},  // '2C99A8'
    {0, 194, 255},   // '00C2FF'
    {52, 69, 147},   // '344593'
    {100, 115, 255}, // '6473FF'
    {0, 24, 236},    // '0018EC'
    {132, 56, 255},  // '8438FF'
    {82, 0, 133},    // '520085'
    {203, 56, 255},  // 'CB38FF'
    {255, 149, 200}, // 'FF95C8'
    {255, 55, 199}   // 'FF37C7'
};






static void dump_tensor_attr(rknn_tensor_attr *attr)
{
    std::string shape_str = attr->n_dims < 1 ? "" : std::to_string(attr->dims[0]);
    for (int i = 1; i < attr->n_dims; ++i)
    {
        shape_str += ", " + std::to_string(attr->dims[i]);
    }

    // printf("  index=%d, name=%s, n_dims=%d, dims=[%s], n_elems=%d, size=%d, w_stride = %d, size_with_stride=%d, fmt=%s, "
    //        "type=%s, qnt_type=%s, "
    //        "zp=%d, scale=%f\n",
    //        attr->index, attr->name, attr->n_dims, shape_str.c_str(), attr->n_elems, attr->size, attr->w_stride,
    //        attr->size_with_stride, get_format_string(attr->fmt), get_type_string(attr->type),
    //        get_qnt_type_string(attr->qnt_type), attr->zp, attr->scale);
}

static unsigned char *load_data(FILE *fp, size_t ofst, size_t sz)
{
    unsigned char *data;
    int ret;

    data = NULL;

    if (NULL == fp)
    {
        return NULL;
    }

    ret = fseek(fp, ofst, SEEK_SET);
    if (ret != 0)
    {
        printf("blob seek failure.\n");
        return NULL;
    }

    data = (unsigned char *)malloc(sz);
    if (data == NULL)
    {
        printf("buffer malloc failure.\n");
        return NULL;
    }
    ret = fread(data, 1, sz, fp);
    return data;
}

static unsigned char *load_model(const char *filename, int *model_size)
{
    FILE *fp;
    unsigned char *data;

    fp = fopen(filename, "rb");
    if (NULL == fp)
    {
        printf("Open file %s failed.\n", filename);
        return NULL;
    }

    fseek(fp, 0, SEEK_END);
    int size = ftell(fp);

    data = load_data(fp, 0, size);

    fclose(fp);

    *model_size = size;
    return data;
}

static int saveFloat(const char *file_name, float *output, int element_size)
{
    FILE *fp;
    fp = fopen(file_name, "w");
    for (int i = 0; i < element_size; i++)
    {
        fprintf(fp, "%.6f\n", output[i]);
    }
    fclose(fp);
    return 0;
}

rkYolov8_seg::rkYolov8_seg(const std::string &model_path)
{
    this->model_path = model_path;
    nms_threshold = NMS_THRESH;      // 默认的NMS阈值
    box_conf_threshold = BOX_THRESH; // 默认的置信度阈值
}

int rkYolov8_seg::init(rknn_context *ctx_in, bool share_weight)
{
    printf("Loading model...\n");
    int model_data_size = 0;
    model_data = load_model(model_path.c_str(), &model_data_size);
    loadLabelName(LABEL_NALE_TXT_PATH,labels);
    // 模型参数复用/Model parameter reuse
    if (share_weight == true)
        ret = rknn_dup_context(ctx_in, &ctx);
    else
        ret = rknn_init(&ctx, model_data, model_data_size, 0, NULL);
    if (ret < 0)
    {
        printf("rknn_init error ret=%d\n", ret);
        return -1;
    }

    // 设置模型绑定的核心/Set the core of the model that needs to be bound
    rknn_core_mask core_mask;
    switch (get_core_num())
    {
    case 0:
        core_mask = RKNN_NPU_CORE_0;
        break;
    case 1:
        core_mask = RKNN_NPU_CORE_1;
        break;
    case 2:
        core_mask = RKNN_NPU_CORE_2;
        break;
    }
    ret = rknn_set_core_mask(ctx, core_mask);
    if (ret < 0)
    {
        printf("rknn_init core error ret=%d\n", ret);
        return -1;
    }

    rknn_sdk_version version;
    ret = rknn_query(ctx, RKNN_QUERY_SDK_VERSION, &version, sizeof(rknn_sdk_version));
    if (ret < 0)
    {
        printf("rknn_init error ret=%d\n", ret);
        return -1;
    }
    printf("sdk version: %s driver version: %s\n", version.api_version, version.drv_version);

    // 获取模型输入输出参数/Obtain the input and output parameters of the model
    ret = rknn_query(ctx, RKNN_QUERY_IN_OUT_NUM, &io_num, sizeof(io_num));
    if (ret < 0)
    {
        printf("rknn_init error ret=%d\n", ret);
        return -1;
    }
    printf("model input num: %d, output num: %d\n", io_num.n_input, io_num.n_output);

    // 设置输入参数/Set the input parameters
    input_attrs = (rknn_tensor_attr *)calloc(io_num.n_input, sizeof(rknn_tensor_attr));
    for (int i = 0; i < io_num.n_input; i++)
    {
        input_attrs[i].index = i;
        ret = rknn_query(ctx, RKNN_QUERY_INPUT_ATTR, &(input_attrs[i]), sizeof(rknn_tensor_attr));
        if (ret < 0)
        {
            printf("rknn_init error ret=%d\n", ret);
            return -1;
        }
        dump_tensor_attr(&(input_attrs[i]));
    }

    // 设置输出参数/Set the output parameters
    output_attrs = (rknn_tensor_attr *)calloc(io_num.n_output, sizeof(rknn_tensor_attr));
    for (int i = 0; i < io_num.n_output; i++)
    {
        output_attrs[i].index = i;
        ret = rknn_query(ctx, RKNN_QUERY_OUTPUT_ATTR, &(output_attrs[i]), sizeof(rknn_tensor_attr));
        dump_tensor_attr(&(output_attrs[i]));
    }


    // Set to context
    rknn_app_ctx.rknn_ctx = ctx;
    // TODO
    if (output_attrs[0].qnt_type == RKNN_TENSOR_QNT_AFFINE_ASYMMETRIC && output_attrs[0].type != RKNN_TENSOR_FLOAT16)
    {
        rknn_app_ctx.is_quant = true;
    }
    else
    {
        rknn_app_ctx.is_quant = false;
    }
    printf("model rknn_app_ctx.is_quant=%d\n",rknn_app_ctx.is_quant);
    rknn_app_ctx.io_num = io_num;
    rknn_app_ctx.input_attrs = (rknn_tensor_attr *)malloc(io_num.n_input * sizeof(rknn_tensor_attr));
    memcpy(rknn_app_ctx.input_attrs, input_attrs, io_num.n_input * sizeof(rknn_tensor_attr));
    rknn_app_ctx.output_attrs = (rknn_tensor_attr *)malloc(io_num.n_output * sizeof(rknn_tensor_attr));
    memcpy(rknn_app_ctx.output_attrs, output_attrs, io_num.n_output * sizeof(rknn_tensor_attr));

    if (input_attrs[0].fmt == RKNN_TENSOR_NCHW)
    {
        printf("model is NCHW input fmt\n");
        channel = input_attrs[0].dims[1];
        height = input_attrs[0].dims[2];
        width = input_attrs[0].dims[3];
    }
    else
    {
        printf("model is NHWC input fmt\n");
        height = input_attrs[0].dims[1];
        width = input_attrs[0].dims[2];
        channel = input_attrs[0].dims[3];
    }
    rknn_app_ctx.model_channel = channel;
    rknn_app_ctx.model_height = height;
    rknn_app_ctx.model_width = width;
    printf("model input height=%d, width=%d, channel=%d\n", height, width, channel);

    memset(inputs, 0, sizeof(inputs));
    inputs[0].index = 0;
    inputs[0].type = RKNN_TENSOR_UINT8;
    inputs[0].size = width * height * channel;
    inputs[0].fmt = RKNN_TENSOR_NHWC;
    inputs[0].pass_through = 0;
    return 0;
}

rknn_context *rkYolov8_seg::get_pctx()
{
    return &ctx;
}
const char *rkYolov8_seg::coco_cls_to_name(int cls_id)
{

    if (cls_id >= OBJ_CLASS_NUM)
    {
        return "null";
    }

    if (labels[cls_id])
    {
        return labels[cls_id];
    }

    return "null";
}
#include "common.h"
int convert_image_with_letterbox(cv::Mat &src_image, cv::Mat &dst_image, letterbox_t* letterbox, char color)
{
    int ret = 0;
    int allow_slight_change = 1;
    int src_w = src_image.cols;
    int src_h = src_image.rows;
    int dst_w = dst_image.cols;
    int dst_h = dst_image.rows;
    int resize_w = dst_w;
    int resize_h = dst_h;

    int padding_w = 0;
    int padding_h = 0;

    int _left_offset = 0;
    int _top_offset = 0;
    float scale = 1.0;

    image_rect_t src_box;
    src_box.left = 0;
    src_box.top = 0;
    src_box.right = src_image.cols - 1;
    src_box.bottom = src_image.rows - 1;

    image_rect_t dst_box;
    dst_box.left = 0;
    dst_box.top = 0;
    dst_box.right = dst_image.cols - 1;
    dst_box.bottom = dst_image.rows - 1;

    float _scale_w = (float)dst_w / src_w;
    float _scale_h = (float)dst_h / src_h;
    if(_scale_w < _scale_h) {
        scale = _scale_w;
        resize_h = (int) src_h*scale;
    } else {
        scale = _scale_h;
        resize_w = (int) src_w*scale;
    }
    // slight change image size for align
    if (allow_slight_change == 1 && (resize_w % 4 != 0)) {
        resize_w -= resize_w % 4;
    }
    if (allow_slight_change == 1 && (resize_h % 2 != 0)) {
        resize_h -= resize_h % 2;
    }
    // padding
    padding_h = dst_h - resize_h;
    padding_w = dst_w - resize_w;
    // center
    if (_scale_w < _scale_h) {
        dst_box.top = padding_h / 2;
        if (dst_box.top % 2 != 0) {
            dst_box.top -= dst_box.top % 2;
            if (dst_box.top < 0) {
                dst_box.top = 0;
            }
        }
        dst_box.bottom = dst_box.top + resize_h - 1;
        _top_offset = dst_box.top;
    } else {
        dst_box.left = padding_w / 2;
        if (dst_box.left % 2 != 0) {
            dst_box.left -= dst_box.left % 2;
            if (dst_box.left < 0) {
                dst_box.left = 0;
            }
        }
        dst_box.right = dst_box.left + resize_w - 1;
        _left_offset = dst_box.left;
    }
    // printf("scale=%f dst_box=(%d %d %d %d) allow_slight_change=%d _left_offset=%d _top_offset=%d padding_w=%d padding_h=%d\n",
    //     scale, dst_box.left, dst_box.top, dst_box.right, dst_box.bottom, allow_slight_change,
    //     _left_offset, _top_offset, padding_w, padding_h);

    //set offset and scale
    if(letterbox != NULL){
        letterbox->scale = scale;
        letterbox->x_pad = _left_offset;
        letterbox->y_pad = _top_offset;
    }
    
    // set rga rect
    im_rect srect={0};
    im_rect drect={0};
    im_rect prect={0};
    rga_buffer_t pat={0};


    srect.x = src_box.left;
    srect.y = src_box.top;
    srect.width = src_box.right - src_box.left + 1;
    srect.height = src_box.bottom - src_box.top + 1;

    drect.x = dst_box.left;
    drect.y = dst_box.top;
    drect.width = dst_box.right - dst_box.left + 1;
    drect.height = dst_box.bottom - dst_box.top + 1;


    if (src_image.type() != CV_8UC3)
    {
        printf("source image type is %d!\n", src_image.type());
        return -1;
    }
    rga_buffer_t rga_buf_dst = wrapbuffer_virtualaddr(dst_image.data, dst_image.cols, dst_image.rows, RK_FORMAT_RGB_888);
    rga_buffer_t rga_buf_src = wrapbuffer_virtualaddr(src_image.data, src_image.cols, src_image.rows, RK_FORMAT_RGB_888);
    // im_rect dst_whole_rect = {0, 0, dst_image->width, dst_image->height};
    // IM_STATUS ret_rga = imfill(rga_buf_dst, dst_whole_rect, 0x72727272);
    // rga process

    IM_STATUS ret_rga = improcess(rga_buf_src, rga_buf_dst, pat, srect, drect, prect, 0);
    return ret;
}
void rkYolov8_seg::drawBoxes(cv::Mat& orig_img, const object_detect_result_list& od_results) {
    char text[256];
    for (int i = 0; i < od_results.count; i++)
    {
        object_detect_result det_result = od_results.results[i];
        sprintf(text, "%s %.1f%%", coco_cls_to_name(det_result.cls_id), det_result.prop * 100);
        // 打印预测物体的信息/Prints information about the predicted object
        // printf("%s @ (%d %d %d %d) %f\n", det_result.name, det_result.box.left, det_result.box.top,
        //        det_result.box.right, det_result.box.bottom, det_result.prop);
        int x1 = det_result.box.left;
        int y1 = det_result.box.top;
        int x2 = det_result.box.right;
        int y2 = det_result.box.bottom;
        // qDebug() << "x1:" << x1 << "y1:" << y1 << "x2:" << x2 << "y2:" << y2 << "prop:" << det_result.prop;
        rectangle(orig_img, cv::Point(x1, y1), cv::Point(x2, y2), cv::Scalar(256, 0, 0, 256), 3);
        putText(orig_img, text, cv::Point(x1, y1 + 12), cv::FONT_HERSHEY_SIMPLEX, 0.4, cv::Scalar(255, 255, 255));
    }
}
void applySegmentationMask(cv::Mat& orig_img, const unsigned char class_colors[][3], object_detect_result_list& od_results) {
    if (od_results.count >= 1)
    {
        int width = orig_img.cols;
        int height = orig_img.rows;
        char *ori_img = (char *)orig_img.data;
        uint8_t *seg_mask = od_results.results_seg[0].seg_mask;
        od_results.results_seg[0].pixel_cnt=0;
        float alpha = 0.5f; // opacity
        for (int j = 0; j < height; j++)
        {
            for (int k = 0; k < width; k++)
            {
                int pixel_offset = 3 * (j * width + k);
                if (seg_mask[j * width + k] != 0)
                {
                    ori_img[pixel_offset + 0] = (unsigned char)clamp(class_colors[seg_mask[j * width + k] % N_CLASS_COLORS][0] * (1 - alpha) + ori_img[pixel_offset + 0] * alpha, 0, 255); // r
                    ori_img[pixel_offset + 1] = (unsigned char)clamp(class_colors[seg_mask[j * width + k] % N_CLASS_COLORS][1] * (1 - alpha) + ori_img[pixel_offset + 1] * alpha, 0, 255); // g
                    ori_img[pixel_offset + 2] = (unsigned char)clamp(class_colors[seg_mask[j * width + k] % N_CLASS_COLORS][2] * (1 - alpha) + ori_img[pixel_offset + 2] * alpha, 0, 255); // b
                    od_results.results_seg[0].pixel_cnt++;
                }
            }
        }
        free(seg_mask);
    }
}
cv::Mat rkYolov8_seg::infer(cv::Mat &orig_img)
{
    //测量执行时间
    // auto total_start = std::chrono::high_resolution_clock::now();
    std::lock_guard<std::mutex> lock(mtx);
    cv::Mat img;
    cv::cvtColor(orig_img, img, cv::COLOR_BGR2RGB);

    //模型输入图像宽度/Model input image width
    int img_width = img.cols;
    int img_height = img.rows;

    letterbox_t letter_box={0};
    cv::Size target_size(width, height);


    // 设置yolov8seg输入图像/Set the input image
    rknn_app_ctx.input_image_width = img_width;
    rknn_app_ctx.input_image_height = img_height;    

    // 图像缩放/Image scaling
    cv::Mat dst_img(target_size.height, target_size.width, CV_8UC3);
    if (img_width != width || img_height != height)
    {
        // rga
        convert_image_with_letterbox(img, dst_img, &letter_box, 0);
        // /*********
        // // opencv
        // float min_scale = std::min(scale_w, scale_h);
        // scale_w = min_scale;
        // scale_h = min_scale;
        // letterbox(img, resized_img, pads, min_scale, target_size);
        // *********/
        inputs[0].buf = dst_img.data;
    }
    else
    {
        inputs[0].buf = img.data;
    }

    rknn_inputs_set(ctx, io_num.n_input, inputs);

    rknn_output outputs[io_num.n_output];
    
    memset(outputs, 0, sizeof(outputs));
    for (int i = 0; i < io_num.n_output; i++)
    {
        outputs[i].want_float = 0;
    }
    //测量执行时间
    // auto start = std::chrono::high_resolution_clock::now();
    // 模型推理/Model inference
    ret = rknn_run(ctx, NULL);
    // qDebug() << "rknn_run=" << ret;
    // auto end = std::chrono::high_resolution_clock::now();
    // qDebug() << "rknn_run => duration:" << std::chrono::duration_cast<std::chrono::microseconds>(end - start).count() << "us";
    ret = rknn_outputs_get(ctx, io_num.n_output, outputs, NULL);
    // qDebug() << "rknn_outputs_get=" << ret;
    // 后处理/Post-processing

    for (int i = 0; i < io_num.n_output; ++i)
    {
        out_scales.push_back(output_attrs[i].scale);
        out_zps.push_back(output_attrs[i].zp);
    }
    // post_process((int8_t *)outputs[0].buf, (int8_t *)outputs[1].buf, (int8_t *)outputs[2].buf, height, width,
    //              box_conf_threshold, nms_threshold, pads, scale_w, scale_h, out_zps, out_scales, &od_results);
#if(1)
    // qDebug() << "output_attrs[0].qnt_type:" << output_attrs[0].qnt_type;
    ret = post_process(&rknn_app_ctx,  outputs, &letter_box, box_conf_threshold,nms_threshold, &od_results);
    // qDebug() << "od_results.count:" << od_results.count;
    // end = std::chrono::high_resolution_clock::now();
    // qDebug() << "rknn_run-->post_process => duration:" << std::chrono::duration_cast<std::chrono::microseconds>(end - start).count() << "us";
#endif
    // //测量执行时间
    // start = std::chrono::high_resolution_clock::now();
    // draw mask
    applySegmentationMask(orig_img, class_colors, od_results);
    // 绘制框体/Draw the box
    drawBoxes(orig_img,od_results);
    //  end = std::chrono::high_resolution_clock::now();
    // qDebug() << "draw mask-->drawBoxes => duration:" << std::chrono::duration_cast<std::chrono::microseconds>(end - start).count() << "us";//3000us

    ret = rknn_outputs_release(ctx, io_num.n_output, outputs);
    // auto total_end = std::chrono::high_resolution_clock::now();
    // cv::imshow("Camera source", orig_img);
    // qDebug() << "total duration:" << std::chrono::duration_cast<std::chrono::microseconds>(total_end - total_start).count() << "us";


    return orig_img;
}

rkYolov8_seg::~rkYolov8_seg()
{
    ret = rknn_destroy(ctx);

    if (model_data)
        free(model_data);

    if (input_attrs)
        free(input_attrs);
    if (output_attrs)
        free(output_attrs);
}
